from langchain_zhipu import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.utilities import SQLDatabase
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

output_parser = StrOutputParser()

llm = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                  , model="glm-4")
llm.temperature = 0.01  # 温度设置为0，结果随机性 ghbnm
mysql_uri = 'mysql+pymysql://rdsroot:Geely%40db20211206@10.240.53.162:3306/test'

template = """Just answer SQL,exclude all information outside of SQL,Based on the table schema below, write a SQL query that would answer the user's question:
{schema}

Question: {question}
SQL Query:"""
prompt = ChatPromptTemplate.from_template(template)
db = SQLDatabase.from_uri(mysql_uri)


def get_schema(_):
    return db.get_table_info()


# def run_query(query):
#     return db.run(query)


sql_response = (
        RunnablePassthrough.assign(schema=get_schema)
        | prompt
        | llm.bind(stop=["\nSQLResult:"])
        | output_parser
)


template = """Based on the table schema below, question, sql query, and sql response, write a natural language response:
{schema}

Question: {question}
SQL Query: {query}
SQL Response: {response}"""
prompt_response = ChatPromptTemplate.from_template(template)

full_chain = (
        RunnablePassthrough.assign(query=sql_response).assign(
            schema=get_schema,
            response=lambda x: db.run(x["query"]),
        )
        | prompt_response
        | llm | output_parser
)

# questionStr = "哪一个学生最好"
# print(sql_response.invoke({"question": questionStr}))


# print(full_chain.invoke({"question": questionStr}))
# print(full_chain.invoke({"question": "which student is best?"}))
print(full_chain.invoke({"question": "How many students are there?"}))